
from numpy import log10, abs, pi, array, zeros
import scipy
from scipy import signal
import matplotlib
import matplotlib.pyplot
import matplotlib as mpl

class Hb:
    def __init__(self,N,com=True):
        self.com = com
        self.N = N
        self.b = signal.firwin(N+1, 0.5)
        self.b[abs(self.b)<1e-4] = 0.0
        a = self.b[self.N//2]
        self.c = self.b[1:self.N//2:2]/a
        self.hN = len(self.c)
        bl = self.N*2
        if com:
            self.buf = zeros(bl,dtype=complex)
        else:
            self.buf = zeros(bl)
    
    def dump(self):
        print('          firwin       ')
        print('------------------------------------')
        for ii in range(self.N+1):
            print(' tap %2d   %-3.6f' % (ii, self.b[ii]))

    def toGo(self):
        a = self.b[self.N//2]
        c = self.b[1:self.N//2:2]/a
        print("[]float64{",end=" ")
        for f in c:
            print(f,end=",")
        print("}")

    def plot(self):
        (wb, Hb) = signal.freqz(self.b)
        fig = mpl.pyplot.figure()
        ax1 = fig.add_subplot(111)
        ax1.plot(wb, 20*log10(abs(Hb)))
        ax1.legend(['firwin'])
        bands = array([0.22,0.28])
        bx = bands*2*pi
        ax1.axvspan(bx[0], bx[1], facecolor='0.5', alpha=0.33)
        ax1.plot(pi/2, -6, 'go')
        ax1.axvline(pi/2, color='g', linestyle='--')
        ax1.axis([0,pi,-64,3])
        ax1.grid('on')
        ax1.set_ylabel('Magnitude (dB)')
        ax1.set_xlabel('Normalized Frequency (radians)')
        ax1.set_title('Half Band Filter Frequency Response')
        mpl.pyplot.show()

    def tail(self,din):
        self.buf[:self.N] = din[-self.N:]
    
    def head(self,din,dout):
        self.buf[self.N:] = din[:self.N]
        n = self.done(self.buf,dout,0)
        if n!=self.N//2:
            print(n,"!=",self.N//2)
    
    def Done(self,din):
        l = len(din)//2
        if self.com:
            dout = zeros(l,dtype=complex)
        else:
            dout = zeros(l)
        self.head(din,dout)
        n = self.done(din,dout,self.N//2)
        if n!=l:
            print(n,"!=",l)
        self.tail(din)
        return dout

    def done(self,din,dout,start):
        l = (len(din)-self.N)//2
        for i in range(l):
            p = (i+self.hN)*2
            v = din[p]
            for j in range(self.hN):
                # print(p+2*(j-self.hN)+1,p,p+2*(self.hN-j)-1)
                v += self.c[j]*(din[p+2*(j-self.hN)+1]+din[p+2*(self.hN-j)-1])
            dout[i+start] = v/2
        return l+start